#!/bin/bash

# Docker镜像导出导入工具
# 支持导出所有本地镜像或导入镜像文件

# 定义颜色变量（避免在某些终端出现乱码）
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 显示彩色信息的函数
print_info() {
    echo -e "${GREEN}[信息]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[警告]${NC} $1"
}

print_error() {
    echo -e "${RED}[错误]${NC} $1"
}

# 导出所有Docker镜像到指定目录
export_images() {
    print_info "正在导出所有Docker镜像..."
    mkdir -p ./docker_images
    success_count=0
    fail_count=0
    
    # 使用while循环处理包含空格的镜像名
    docker images --format "{{.Repository}}:{{.Tag}}" | while IFS= read -r img; do
        # 跳过<none>镜像
        case "$img" in
            *"<none>"*) continue ;;
        esac

        # 替换镜像名中的特殊字符以创建合法的文件名
        filename=$(echo "$img" | sed 's/\//_/g; s/:/_/g')

        print_info "正在导出镜像: $img"
        if docker save "$img" -o "./docker_images/${filename}.tar"; then
            success_count=$((success_count + 1))
            print_info "成功导出: $img"
        else
            fail_count=$((fail_count + 1))
            print_error "导出失败: $img"
        fi
    done
    
    print_info "镜像导出完成。成功: $success_count, 失败: $fail_count"
    print_info "所有镜像已导出到 ./docker_images 目录"
}

# 导入当前目录下所有镜像文件
import_images() {
    print_info "正在导入当前目录下的所有镜像文件..."
    local success_count=0
    local fail_count=0
    
    if [ ! -d "./docker_images" ]; then
        print_error "未找到 ./docker_images 目录，请确认目录存在"
        return 1
    fi
    
    for img in ./docker_images/*.tar; do
        if [ -f "$img" ]; then
            print_info "正在导入镜像文件: $(basename "$img")"
            if docker load -i "$img"; then
                success_count=$((success_count + 1))
                print_info "成功导入: $(basename "$img")"
            else
                fail_count=$((fail_count + 1))
                print_error "导入失败: $(basename "$img")"
            fi
        fi
    done
    
    print_info "镜像导入完成。成功: $success_count, 失败: $fail_count"
}

# 单个镜像导出功能
export_single_image() {
    printf "请输入要导出的镜像ID或名称 (例如 nginx:latest): "
    read image_id
    if [ -z "$image_id" ]; then
        print_error "镜像ID不能为空"
        return 1
    fi

    printf "请输入导出文件名 (例如 nginx.tar): "
    read export_path
    if [ -z "$export_path" ]; then
        print_error "导出文件名不能为空"
        return 1
    fi
    
    print_info "开始导出镜像 $image_id 到 $export_path"
    if docker save "$image_id" -o "$export_path"; then
        print_info "镜像导出完成: $export_path"
    else
        print_error "镜像导出失败"
        return 1
    fi
}

# 单个镜像导入功能
import_single_image() {
    printf "请输入导入文件路径 (例如 ./nginx.tar): "
    read import_path
    if [ ! -f "$import_path" ]; then
        print_error "文件不存在: $import_path"
        return 1
    fi
    
    print_info "开始导入镜像文件: $import_path"
    if docker load -i "$import_path"; then
        print_info "镜像导入完成"
    else
        print_error "镜像导入失败"
        return 1
    fi
}

# 主菜单
main_menu() {
    echo ""
    echo "=============================================="
    echo "         Docker 镜像管理工具"
    echo "=============================================="
    echo "请选择操作："
    echo "1. 导出所有本地Docker镜像"
    echo "2. 导入当前目录下的所有镜像文件"
    echo "3. 导出单个Docker镜像"
    echo "4. 导入单个镜像文件"
    echo "0. 退出程序"
    echo "=============================================="
}

# 主程序循环
while true; do
    main_menu
    printf "请输入选项 (0-4): "
    read choice

    case $choice in
        1)
            export_images
            ;;
        2)
            import_images
            ;;
        3)
            export_single_image
            ;;
        4)
            import_single_image
            ;;
        0)
            print_info "感谢使用，再见！"
            exit 0
            ;;
        *)
            print_error "无效的选项，请输入 0-4 之间的数字"
            ;;
    esac
    
    echo ""
    printf "按回车键继续..."
    read dummy
    
    # 添加一个空行确保循环正常执行
    echo ""
    
    # 确保循环可以正常继续
    continue
    
    # 避免语法错误
    
    # 如果上面的代码仍然有问题，这里添加一个简单的结束语句
    
    # 这个位置应该不会被执行，因为前面已经有continue
    
    # 为了确保语法正确，我们在这里添加一个空的命令
    :
done